#include "encoder.h"
#include "bootstrapping/polyarith.h"

Encoder::Encoder(Context *params)
{
    this->params = params;

    uint64_t m, pos, index1, index2;
    uint64_t slots = params->ringT->N; // T的长度为槽数
    indexMatrix.resize(slots);
    auto logN = uint64_t(log2(double(slots)));
    uint64_t rowSize = slots >> 1;
    m = uint64_t(slots) << 1;
    pos = 1;

    for (int i = 0; i < rowSize; i++) {
        index1 = (pos - 1) >> 1;
        index2 = (m - pos - 1) >> 1;

        indexMatrix[i] = BitReverse64(index1, logN);
        indexMatrix[i | rowSize] = BitReverse64(index2, logN);

        pos *= GaloisGen;
        pos &= (m - 1);
    }

    rescaleParams.resize(params->ringQ->Modulus.size());
    for (int i = 0; i < params->ringQ->Modulus.size(); i++) {
        uint64_t qi = params->ringQ->Modulus[i];
        rescaleParams[i] = InvMod_(params->T(), qi), qi, params->ringQ->bredParams[i];
    }

    scaler = new RNSScaler(params->ringQ, params->ringT);
    tmpPoly = params->ringT->NewPoly();
    tmpPtRt = new PlaintextRingT(params, 0);
}
// 拷贝构造函数的实现
Encoder::Encoder(const Encoder &other)
    : indexMatrix(other.indexMatrix),
      rescaleParams(other.rescaleParams),
      params(nullptr),
      scaler(nullptr),
      tmpPoly(nullptr),
      tmpPtRt(nullptr)
{
    // 深拷贝指针成员变量
    if (other.params) {
        this->params = new Context(*other.params);
    }

    if (other.scaler) {
        this->scaler = new RNSScaler(*other.scaler);
    }

    if (other.tmpPoly) {
        this->tmpPoly = new Poly(*other.tmpPoly);
    }

    if (other.tmpPtRt) {
        this->tmpPtRt = new PlaintextRingT(*other.tmpPtRt);
    }
}
Encoder &Encoder::operator = (const Encoder &other)
{
    // 防止自赋值
    if (this == &other)
        return *this;

    // 复制非指针成员变量
    this->indexMatrix = other.indexMatrix;
    this->rescaleParams = other.rescaleParams;

    // 复制指针成员变量（深拷贝或浅拷贝）
    if (other.params) {
        this->params = new Context(*other.params); // 深拷贝
    } else {
        this->params = nullptr;
    }

    if (other.scaler) {
        this->scaler = new RNSScaler(*other.scaler); // 深拷贝
    } else {
        this->scaler = nullptr;
    }

    if (other.tmpPoly) {
        this->tmpPoly = new Poly(*other.tmpPoly); // 深拷贝
    } else {
        this->tmpPoly = nullptr;
    }

    if (other.tmpPtRt) {
        this->tmpPtRt = new PlaintextRingT(*other.tmpPtRt); // 深拷贝
    } else {
        this->tmpPtRt = nullptr;
    }

    return *this;
}
Encoder::~Encoder()
{
    delete this->scaler;
    delete this->tmpPoly;
    delete this->tmpPtRt;
    delete this->params;
}
/*
 * @brief 将一个无符号整数系数向量编码为 RingT 类型的明文。
 *
 * 该函数将输入的无符号整数系数 `coeffs` 编码到 `PlaintextRingT` 类型的明文对象 `pt` 中。编码过程中会根据指定的
 * `indexMatrix` 将系数映射到 `pt` 的正确位置。
 *
 * @param coeffs 输入的无符号整数系数向量，表示要编码的系数。
 * @param pt 指向目标 `PlaintextRingT` 对象的指针，编码后的数据将存储在该对象中。
 */
void Encoder::EncodeUintRingT(const std::vector<uint64_t> &coeffs, PlaintextRingT *pt)
{
    if (coeffs.size() > this->indexMatrix.size()) {
        throw std::invalid_argument(
            "invalid input to encode: number of coefficients must be smaller or equal to the ring degree");
    }
    if (pt->value->Coeffs[0].size() != this->indexMatrix.size()) {
        throw std::invalid_argument(
            "invalid plaintext to receive encoding: number of coefficients does not match the ring degree");
    }
    for (size_t i = 0; i < coeffs.size(); i++) {
        pt->value->Coeffs[0][this->indexMatrix[i]] = coeffs[i];
    }
    for (size_t i = coeffs.size(); i < this->indexMatrix.size(); i++) {
        pt->value->Coeffs[0][this->indexMatrix[i]] = 0;
    }
    this->params->ringT->InvNTT(pt->value, pt->value);
}
/*
 * @brief 将一个无符号整数系数向量编码为 RingQ 类型的明文。
 *
 * 该函数将输入的无符号整数系数 `coeffs` 编码到 `Plaintext` 类型的明文对象 `pt`
 * 中。先将系数编码到RingT上，然后缩放到RingQ上。
 *
 * @param coeffs 输入的无符号整数系数向量，表示要编码的系数。
 * @param pt 指向目标 `Plaintext` 对象的指针，编码后的数据将存储在该对象中。
 */
void Encoder::EncodeUint(const std::vector<uint64_t> &coeffs, Plaintext *pt)
{
    PlaintextRingT ptRt(this->params, this->params->QCount() - 1);
    ptRt.value = pt->value;
    EncodeUintRingT(coeffs, &ptRt);
    ScaleUp(&ptRt, pt);
}
/*
 * @brief 对明文进行缩放操作。
 *
 * 该函数将对传入的 `PlaintextRingT` 和 `Plaintext` 类型的明文对象进行缩放操作，将明文从环T转换到环Q上。
 *
 * @param ptRt 指向 `PlaintextRingT` 类型对象的指针，表示将要被缩放的明文数据。
 * @param pt 指向 `Plaintext` 类型对象的指针，表示输入的明文数据，进行缩放后的结果将存储在该对象中。
 */
void Encoder::ScaleUp(PlaintextRingT *ptRt, Plaintext *pt)
{
    ScaleUp(this->params->ringQ, this->params->ringT, this->tmpPoly->Coeffs[0], ptRt->value, pt->value);
}
/*
 * @brief 对明文数据进行缩放操作
 *
 * 该函数对传入的多项式 `pIn` 进行缩放处理，并将结果存储在 `pOut` 中。
 *
 * @param ringQ 指向 `Ring` 类型对象的指针，表示与输入多项式的模数相关的环。
 * @param ringT 指向 `Ring` 类型对象的指针，表示目标环，用于缩放操作中的模数。
 * @param tmp 一个 `std::vector<uint64_t>` 类型的引用，用于存储临时计算结果。
 * @param pIn 指向输入多项式的指针，包含要进行缩放的系数。
 * @param pOut 指向输出多项式的指针，用于存储缩放后的结果。
 */
void Encoder::ScaleUp(Ring *ringQ, Ring *ringT, std::vector<uint64_t> &tmp, Poly *pIn, Poly *pOut)
{
    mpz_class remainder = *(ringQ->modulusBigint) % *(ringT->modulusBigint);
    uint64_t qModTmontgomery = remainder.get_ui();

    uint64_t t = ringT->Modulus[0];
    uint64_t tHalf = t >> 1;
    vector<uint64_t> mBredParams = ringT->bredParams[0];
    // (x * Q + T/2) mod T
#pragma omp parallel for
    for (int i = 0; i < ringQ->N; i++) {
        MulModBarrett(tmp[i], pIn->Coeffs[0][i], qModTmontgomery, t, mBredParams[2], mBredParams[0], mBredParams[1]);
        tmp[i] = CRed(tmp[i] + tHalf, t); // 计算pIn + t/2 再模t
    }
#pragma omp parallel for
    // (x * T^-1 - T/2) mod Qi
    for (int i = 0; i < pOut->Coeffs.size(); i++) {
        std::vector<uint64_t> &p1tmp = pOut->Coeffs[i];
        uint64_t qi = ringQ->Modulus[i];
        std::vector<uint64_t> bredParams = ringQ->bredParams[i];
        uint64_t rescaleParams = qi - this->rescaleParams[i]; // rescaleParams[i]: T^(qi-2) mod qi
        uint64_t tHalfNegQi = qi - BRedAdd(tHalf, qi, bredParams);

        for (int j = 0; j < ringQ->N; j++) {
            MulModBarrett(p1tmp[j], tmp[j] + tHalfNegQi, rescaleParams, qi, bredParams[2], bredParams[0],
                bredParams[1]);
        }
    }
}
/*
 * @brief 对明文数据进行缩放操作，将RingQ上的明文转换到RingT上
 *
 * 该函数将对输入的 `Plaintext` 对象 `pt` 和 `PlaintextRingT` 对象 `ptRt` 进行缩放操作。
 * @param pt 指向输入的 `Plaintext` 类型对象的指针，包含要进行缩放的明文数据。
 * @param ptRt 指向输出的 `PlaintextRingT` 类型对象的指针，缩放后的结果将存储在该对象中。
 */
void Encoder::ScaleDown(Plaintext *pt, PlaintextRingT *ptRt) const
{
    this->scaler->DivByQOverTRounded(pt->value, ptRt->value);
}

/*
 * @brief 解码并将明文数据转换为 `PlaintextRingT` 类型。
 *
 * 该函数尝试将输入的明文数据（`Plaintext` 类型）解码为 `PlaintextRingT` 类型的数据。
 *
 * @param p 输入的明文数据，指向 `void` 类型的指针，实际类型应为 `Plaintext`。
 * @param ptRt 指向 `PlaintextRingT` 类型的对象，解码后的数据将存储在该对象中。
 */
void Encoder::DecodeRingT(void *p, PlaintextRingT *ptRt)
{
    Plaintext *plaintext = static_cast<Plaintext *>(p);

    if (plaintext != nullptr) {
        ScaleDown(plaintext, ptRt);
        return;
    }
    throw std::invalid_argument("unsupported plaintext type");
}
/*
 * @brief 解码明文数据并将其转换为无符号整数消息向量。
 *
 * 该函数首先将输入的 `Plaintext` 对象解码为 `PlaintextRingT` 类型的数据，随后进行系数映射并将其存储在 `coeffs` 向量中。
 *
 * @param p 输入的明文数据，指向 `void` 类型的指针，实际类型为 `Plaintext`。
 * @param coeffs 用于存储解码后无符号整数系数的 `std::vector<uint64_t>` 引用。
 */
void Encoder::DecodeUint(void *p, std::vector<uint64_t> &coeffs)
{
    PlaintextRingT *ptRt;
    PlaintextRingT *tmpPtRt12 = new PlaintextRingT(this->params, 0);

    p = static_cast<Plaintext *>(p);
    if (p != nullptr) {
        DecodeRingT(p, tmpPtRt12);
        ptRt = tmpPtRt12;
    }
    this->params->ringT->NTT(ptRt->value, tmpPoly);

    for (uint64_t i = 0; i < this->params->ringQ->N; i++) {
        coeffs[i] = tmpPoly->Coeffs[0][indexMatrix[i]];
    }
}

std::vector<uint64_t> Encoder::DecodeUintNew(void *p)
{
    std::vector<uint64_t> coeffs(this->params->ringQ->N, 0);
    DecodeUint(p, coeffs);
    return coeffs;
}

/*
 * @brief 计算给定数值最大可被 2 整除的幂次。
 * @param x 输入的无符号整数。
 * @return 如果 `x` 不为零，返回 `x` 最大可被 2 整除的幂次；如果 `x` 为零，返回 `uint64_t` 类型的位数（通常为 64）。
 */
inline uint64_t HighestDividingPower2(uint64_t x)
{
    if (x == 0) {
        return std::numeric_limits<uint64_t>::digits;
    }
    unsigned long result = 0;
    while ((x & 1) == 0) {
        result += 1;
        x = x >> 1;
    }
    return result;
}
/*
 * @brief 对多项式进行归一化，去除末尾的零系数。
 * @param p 输入和输出的多项式，类型为 `poly`（通常是一个整数类型的向量或数组）。该多项式会被修改，去掉末尾的零系数。
 */
void PolyNormalize(poly &p)
{
    int64_t i = p.size() - 1;
    while (i >= 0 && p[i] == 0) {
        --i;
    }
    p.resize(i + 1);
}
